import numpy as np
import pickle
from geometric_clustering import compute_curvatures, cluster_signed_modularity
import sys

from cdlib.algorithms import louvain, infomap, spinglass, belief
import cdlib
import networkx as nx

from GraphRicciCurvature.OllivierRicci import OllivierRicci
from GraphRicciCurvature.util import cut_graph_by_cutoff
import time


if __name__ == "__main__":
    
    res = pickle.load(open('LFR_benchmark.pkl','rb'))
   
    for i, G in enumerate(res['graphs']):
        # if i<70:
        #     continue
        print(i)
        
        G.remove_edges_from(nx.selfloop_edges(G))
        
        #compute curvatures
        print('Cmputing curvatures')
        times = np.logspace(-.4, .2, 25)
        epsilon = np.finfo(float).eps
        kappas = compute_curvatures(G, times, n_workers=16, use_spectral_gap=False, measure_cutoff=epsilon)
        
        #geometric modularity clustering
        print('Running geometric modularity')
        geometric_modularity = cluster_signed_modularity(
                G,
                times,
                kappas,
                kappa0=None,
                n_louvain=100,
                n_louvain_VI=100,
                n_workers=16,
                with_postprocessing=True,
        )
        res['geometric_modularity'][i] = geometric_modularity
        
        #modularity
        print('Running modularity')
        modularity = louvain(G)
        
        # #run spinglass
        print('Running spinglass')
        sg= spinglass(G)
        
        print('Running belief propagation')
        belief_propagation = belief(G, q_max=40)
        
        #infomap
        print('Running infomap')
        im = infomap(G)
        
        #Ricci flow
        print('Running Ricci-flow')
        ricci_flow = OllivierRicci(G,alpha=0.5)
        ricci_flow.compute_ricci_flow(iterations=50)
        try:
            rf = ricci_flow.ricci_community()
            rf = rf[1]
        except:
            n_comms = len(res['ground_truth'][i].communities)
            n = len(res['graphs'][i].nodes)
            rf = np.random.choice([i for i in range(n_comms)],n)
            
        rf = [[n for n in G.nodes if rf[n]==l] for l in range(max(rf)+1)]
        rf = cdlib.NodeClustering(communities=rf, graph=G, method_name="reference")
        
        #collect metadata and results
        if 'curvatures' not in res.keys():
            res['curvatures'] = []
        res['curvatures'].append(kappas)
        
        if 'geometric_modularity' not in res.keys():
            res['geometric_modularity'] = []
        res['geometric_modularity'].append(geometric_modularity)
        
        if 'belief_propagation' not in res.keys():
            res['belief_propagation'] = []
        res['belief_propagation'].append(belief_propagation)
        
        if 'modularity' not in res.keys():
            res['modularity'] = []
        res['modularity'].append(modularity)
        
        if 'spinglass' not in res.keys():
            res['spinglass'] = []
        res['spinglass'].append(sg)
        
        if 'infomap' not in res.keys():
            res['infomap'] = []
        res['infomap'].append(im)
        
        if 'ricciflow' not in res.keys():
            res['ricciflow'] = []
        res['ricciflow'].append(rf)
        
        #saving
        pickle.dump(res, open('LFR_benchmark.pkl','wb'))
